/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <map>
#include "bsr_common.h"
#include "parbsr_ij_generator.h"

namespace YHAMG
{

ParBSR_IJ_Generator::ParBSR_IJ_Generator()
	: Loc_rows(0),
	block_size(0),
	col_partitioning(0),
	LocGlo_rowptr(0),
	LocGlo_colind(0),
	LocGlo_values(0),
	pattern_symmetry(0)
{
}

ParBSR_IJ_Generator::ParBSR_IJ_Generator(MPI_Comm comm)
	: ParBSR_Generator(comm),
	Loc_rows(0),
	block_size(0),
	col_partitioning(0),
	LocGlo_rowptr(0),
	LocGlo_colind(0),
	LocGlo_values(0),
	pattern_symmetry(0)
{
}

ParBSR_IJ_Generator::ParBSR_IJ_Generator(MPI_Comm comm, int _Loc_rows, int _block_size, const long* _col_partitioning, const int* _LocGlo_rowptr, const long* _LocGlo_colind, const double* _LocGlo_values, int _pattern_symmetry)
	: ParBSR_Generator(comm),
	Loc_rows(_Loc_rows),
	block_size(_block_size),
	col_partitioning(_col_partitioning),
	LocGlo_rowptr(_LocGlo_rowptr),
	LocGlo_colind(_LocGlo_colind),
	LocGlo_values(_LocGlo_values),
	pattern_symmetry(_pattern_symmetry)
{
}

void ParBSR_IJ_Generator::operator()(ParBSR_Matrix& A) const
{
	int MPI_SIZE;
	int MPI_RANK;
	MPI_Comm_size(comm, &MPI_SIZE);
	MPI_Comm_rank(comm, &MPI_RANK);

	long starting_col = col_partitioning[MPI_RANK];
	long ending_col = col_partitioning[MPI_RANK + 1];

	int block_nnz = block_size * block_size;

	int* LocLoc_rowptr = new int[Loc_rows + 1];
	int* LocExt_rowptr = new int[Loc_rows + 1];

	char* recv_nb = new char[MPI_SIZE];
	for (int r = 0; r < MPI_SIZE; ++r)
		recv_nb[r] = 0;

	LocLoc_rowptr[0] = 0;
	LocExt_rowptr[0] = 0;

	for (int i = 0, cnt_loc = 0, cnt_ext = 0; i < Loc_rows; LocLoc_rowptr[++i] = cnt_loc, LocExt_rowptr[i] = cnt_ext)
	{
		for (int j = LocGlo_rowptr[i]; j < LocGlo_rowptr[i + 1]; ++j)
			++(LocGlo_colind[j] >= starting_col && LocGlo_colind[j] < ending_col ? cnt_loc : cnt_ext);
	}

	int* LocLoc_colind = new int[LocLoc_rowptr[Loc_rows]];
	double* LocLoc_values = new double[LocLoc_rowptr[Loc_rows] * block_nnz];
	int* LocExt_colind = new int[LocExt_rowptr[Loc_rows]];
	int* LocExt_colfrom = new int[LocExt_rowptr[Loc_rows]];
	double* LocExt_values = new double[LocExt_rowptr[Loc_rows] * block_nnz];

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < Loc_rows; ++i)
	{
		for (int j = LocGlo_rowptr[i], k = LocLoc_rowptr[i], t = LocExt_rowptr[i]; j < LocGlo_rowptr[i + 1]; ++j)
		{
			long jcol = LocGlo_colind[j];

			if (jcol < starting_col)
			{
				for (int r = MPI_RANK - 1; r >= 0; --r)
				{
					if (jcol >= col_partitioning[r])
					{
						LocExt_colind[t] = (int)(jcol - col_partitioning[r]);
						LocExt_colfrom[t] = r;
						BSRBlockCopy(block_size, LocGlo_values + j * block_nnz, LocExt_values + t * block_nnz);
						++t;
						recv_nb[r] = 1;
						break;
					}
				}
			}
			else if (jcol >= ending_col)
			{
				for (int r = MPI_RANK + 1; r < MPI_SIZE; ++r)
				{
					if (jcol < col_partitioning[r + 1])
					{
						LocExt_colind[t] = (int)(jcol - col_partitioning[r]);
						LocExt_colfrom[t] = r;
						BSRBlockCopy(block_size, LocGlo_values + j * block_nnz, LocExt_values + t * block_nnz);
						++t;
						recv_nb[r] = 1;
						break;
					}
				}
			}
			else
			{
				LocLoc_colind[k] = (int)(jcol - starting_col);
				BSRBlockCopy(block_size, LocGlo_values + j * block_nnz, LocLoc_values + k * block_nnz);
				++k;
			}
		}
	}

	int num_neighbors = 0;
	std::map<int, int> map_neighbor;

	if (pattern_symmetry)
	{
		for (int r = 0; r < MPI_SIZE; ++r)
			if (recv_nb[r]) map_neighbor.insert(std::pair<int, int>(r, num_neighbors++));
	}
	else
	{
		char* send_nb = new char[MPI_SIZE];

		MPI_Alltoall(recv_nb, 1, MPI_CHAR, send_nb, 1, MPI_CHAR, comm);

		for (int r = 0; r < MPI_SIZE; ++r)
			if (recv_nb[r] || send_nb[r]) map_neighbor.insert(std::pair<int, int>(r, num_neighbors++));

		delete[] send_nb;
	}

	std::map<int, int>* map_Ext = new std::map<int, int>[num_neighbors];
	for (int i = 0; i < Loc_rows; ++i)
	{
		for (int j = LocExt_rowptr[i]; j < LocExt_rowptr[i + 1]; ++j)
		{
			int jfrom = map_neighbor[LocExt_colfrom[j]];
			LocExt_colfrom[j] = jfrom;
			std::pair<std::map<int, int>::iterator, bool> ret = map_Ext[jfrom].insert(std::pair<int, int>(LocExt_colind[j], map_Ext[jfrom].size()));
			LocExt_colind[j] = ret.first->second;
		}
	}

	int* neighbor_ranks = new int[num_neighbors];
	for (std::map<int, int>::iterator iter = map_neighbor.begin(); iter != map_neighbor.end(); ++iter)
		neighbor_ranks[iter->second] = iter->first;

	int* Ext_ptr = new int[num_neighbors + 1];
	Ext_ptr[0] = 0;
	for (int r = 0; r < num_neighbors; ++r)
		Ext_ptr[r + 1] = Ext_ptr[r] + map_Ext[r].size();

	int* Ext_map = new int[Ext_ptr[num_neighbors]];
	for (int r = 0; r < num_neighbors; ++r)
		for (std::map<int, int>::iterator iter = map_Ext[r].begin(); iter != map_Ext[r].end(); ++iter)
			Ext_map[Ext_ptr[r] + iter->second] = iter->first;

	for (int j = 0; j < LocExt_rowptr[Loc_rows]; ++j)
		LocExt_colind[j] += Ext_ptr[LocExt_colfrom[j]];

	delete[] map_Ext;
	delete[] LocExt_colfrom;
	delete[] recv_nb;

	A.Free();

	A.comm = comm;

	A.LocLoc.size[0] = Loc_rows;
	A.LocLoc.size[1] = (int)(ending_col - starting_col);
	A.LocLoc.block_size = block_size;
	A.LocLoc.rowptr = LocLoc_rowptr;
	A.LocLoc.colind = LocLoc_colind;
	A.LocLoc.values = LocLoc_values;

	A.LocExt.size[0] = Loc_rows;
	A.LocExt.size[1] = Ext_ptr[num_neighbors];
	A.LocExt.block_size = block_size;
	A.LocExt.rowptr = LocExt_rowptr;
	A.LocExt.colind = LocExt_colind;
	A.LocExt.values = LocExt_values;

	A.num_neighbors = num_neighbors;
	A.neighbor_ranks = neighbor_ranks;

	A.Ext_ptr = Ext_ptr;
	A.Ext_map = Ext_map;
}

}